package middleware

import (
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"sdapiproject/models"
	"sdapiproject/redis_func"
	"sdapiproject/services"
	"sdapiproject/utils"
	"strconv"
)

func Txt2imgTask(prompt, TaskId string, c *gin.Context) {

	fmt.Println("进到Txt2imgTask的任务")

	status, _ := redis_func.GetTxt2ImgStatusFromRedis(TaskId)
	if status == redis_func.CANCEL {
		fmt.Println("任务已取消")
		return
	}

	redis_func.SetTxt2ImgStatusToRedis(TaskId, redis_func.PROCESSING)
	url := utils.Cfg.SDWebUI.Url + "/sdapi/v1/txt2img"

	data := models.NewData()

	prompt, err := utils.TranslateToENOne(prompt)
	if err != nil {
		fmt.Println("翻译失败:", err)
		return
	}
	data.Prompt = prompt
	data.TaskId = TaskId

	jsonData, err := json.Marshal(data)
	if err != nil {
		redis_func.SetTxt2ImgStatusToRedis(data.TaskId, redis_func.FAILED)
		c.JSON(http.StatusInternalServerError, utils.ErrorResponse(http.StatusInternalServerError, err.Error()))
		return
	}

	resp, err := services.SubmitPost(url, jsonData)
	if err != nil {
		redis_func.SetTxt2ImgStatusToRedis(TaskId, redis_func.FAILED)
		fmt.Println("请求失败")
		return
	}

	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		redis_func.SetTxt2ImgStatusToRedis(TaskId, redis_func.FAILED)
		log.Println(err)
		return
	}

	var response models.SuccessResponse

	err = json.Unmarshal(body, &response)
	if err != nil {
		redis_func.SetTxt2ImgStatusToRedis(TaskId, redis_func.FAILED)
		log.Println(err)
		return
	}

	dir, err := utils.CreateTaskDir(TaskId)
	if err != nil || dir == "" {
		redis_func.SetTxt2ImgStatusToRedis(TaskId, redis_func.FAILED)
		log.Println(err)
		return
	}

	for i, img := range response.Images {

		name := filepath.Join(dir, "img"+strconv.Itoa(i+1)+".png")
		err = utils.SaveEncodedImage(img, name)
		if err != nil {
			log.Println(err)
			return
		}
		response.Images[i] = name
	}

	result, _ := json.Marshal(response)

	err = redis_func.SetTxt2ImgTaskToRedis(TaskId, string(result))
	if err != nil {
		log.Println(err)
		return
	}
	redis_func.SetTxt2ImgStatusToRedis(TaskId, redis_func.SUCCESS)

}
